package com.ai.model.xunfei;

import com.ai.model.ChatRequest;
import okhttp3.HttpUrl;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.WebSocket;
import org.springframework.stereotype.Component;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.text.SimpleDateFormat;
import java.util.*;

/**
 * 星火大模型调用
 *
 * @author hh
 * @since 2024/5/26
 */
@Component
public class XfChatAiRequest implements ChatRequest {

    private final XfApiProperties xfApiProperties;

    public XfChatAiRequest(XfApiProperties xfApiProperties) {
        this.xfApiProperties = xfApiProperties;
    }

    public String chat(String newQuestion) throws Exception {
        String authUrl = getAuthUrl();

        OkHttpClient okHttpClient = new OkHttpClient().newBuilder().build();
        Request request = new Request.Builder().url(authUrl).build();
        XfChatAiWebsocket xfChatAiWebsocket = new XfChatAiWebsocket(xfApiProperties, newQuestion);
        WebSocket webSocket = okHttpClient.newWebSocket(request, xfChatAiWebsocket);
        do {
            try {
                Thread.sleep(200);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        } while (!xfChatAiWebsocket.isFinished());
        webSocket.close(1000, "");
        return xfChatAiWebsocket.getAnswer();
    }

    private String getAuthUrl() throws NoSuchAlgorithmException, MalformedURLException, InvalidKeyException {
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(xfApiProperties.getApiSecret().getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        URL url = new URL(xfApiProperties.getHostUrl());
        String date = getDate();
        String preStr = "host: " + url.getHost() + "\n" +
                "date: " + date + "\n" +
                "GET " + url.getPath() + " HTTP/1.1";
        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        String sha = Base64.getEncoder().encodeToString(hexDigits);

        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", xfApiProperties.getApiKey(), "hmac-sha256", "host date request-line", sha);

        HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath())).newBuilder()
                .addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8)))
                .addQueryParameter("date", date)
                .addQueryParameter("host", url.getHost())
                .build();
        return httpUrl.toString().replace("http://", "ws://").replace("https://", "wss://");
    }

    private String getDate() {
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        return format.format(new Date());
    }
}
